package com.baojie.zk.example.watcher.cloud;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.springframework.context.ApplicationContext;
import org.springframework.context.ApplicationContextInitializer;
import org.springframework.context.ConfigurableApplicationContext;
import org.springframework.core.Ordered;
import org.springframework.core.env.*;
import org.springframework.security.crypto.encrypt.TextEncryptor;

import java.util.*;
import java.util.regex.Pattern;

public class EnvironmentDecryptApplicationInitializer implements ApplicationContextInitializer<ConfigurableApplicationContext>, Ordered {

    public static final String DECRYPTED_PROPERTY_SOURCE_NAME = "decrypted";

    public static final String DECRYPTED_BOOTSTRAP_PROPERTY_SOURCE_NAME = "decryptedBootstrap";

    private int order = Ordered.HIGHEST_PRECEDENCE + 15;

    private static Log logger = LogFactory.getLog(EnvironmentDecryptApplicationInitializer.class);

    private TextEncryptor encryptor;

    private boolean failOnError = true;

    /**
     * Strategy to determine how to handle exceptions during decryption.
     *
     * @param failOnError The flag value (default true).
     */
    public void setFailOnError(boolean failOnError) {
        this.failOnError = failOnError;
    }

    public EnvironmentDecryptApplicationInitializer(TextEncryptor encryptor) {
        this.encryptor = encryptor;
    }

    @Override
    public int getOrder() {
        return this.order;
    }

    public void setOrder(int order) {
        this.order = order;
    }

    @Override
    public void initialize(ConfigurableApplicationContext applicationContext) {

        ConfigurableEnvironment environment = applicationContext.getEnvironment();
        MutablePropertySources propertySources = environment.getPropertySources();

        Set<String> found = new LinkedHashSet<>();
        Map<String, Object> map = decrypt(propertySources);
        if (!map.isEmpty()) {
            // We have some decrypted properties
            found.addAll(map.keySet());
            insert(applicationContext, new SystemEnvironmentPropertySource(
                    DECRYPTED_PROPERTY_SOURCE_NAME, map));
        }
        PropertySource<?> bootstrap = propertySources.get(BootstrapApplicationListener.BOOTSTRAP_PROPERTY_SOURCE_NAME);
        if (bootstrap != null) {
            map = decrypt(bootstrap);
            if (!map.isEmpty()) {
                found.addAll(map.keySet());
                insert(applicationContext, new SystemEnvironmentPropertySource(
                        DECRYPTED_BOOTSTRAP_PROPERTY_SOURCE_NAME, map));
            }
        }
        if (!found.isEmpty()) {
            ApplicationContext parent = applicationContext.getParent();
            if (parent != null) {
                // The parent is actually the bootstrap context, and it is fully
                // initialized, so we can fire an EnvironmentChangeEvent there to rebind
                // @ConfigurationProperties, in case they were encrypted.
                parent.publishEvent(new EnvironmentChangeEvent(parent, found));
            }

        }
    }

    private void insert(ApplicationContext applicationContext,
            PropertySource<?> propertySource) {
        ApplicationContext parent = applicationContext;
        while (parent != null) {
            if (parent.getEnvironment() instanceof ConfigurableEnvironment) {
                ConfigurableEnvironment mutable = (ConfigurableEnvironment) parent
                        .getEnvironment();
                insert(mutable.getPropertySources(), propertySource);
            }
            parent = parent.getParent();
        }
    }

    private void insert(MutablePropertySources propertySources, PropertySource<?> propertySource) {
        if (propertySources.contains(BootstrapApplicationListener.BOOTSTRAP_PROPERTY_SOURCE_NAME)) {
            if (DECRYPTED_BOOTSTRAP_PROPERTY_SOURCE_NAME
                    .equals(propertySource.getName())) {
                propertySources.addBefore(BootstrapApplicationListener.BOOTSTRAP_PROPERTY_SOURCE_NAME, propertySource);
            } else {
                propertySources.addAfter(BootstrapApplicationListener.BOOTSTRAP_PROPERTY_SOURCE_NAME, propertySource);
            }
        } else {
            propertySources.addFirst(propertySource);
        }
    }

    public Map<String, Object> decrypt(PropertySources propertySources) {
        Map<String, Object> overrides = new LinkedHashMap<>();
        List<PropertySource<?>> sources = new ArrayList<>();
        for (PropertySource<?> source : propertySources) {
            sources.add(0, source);
        }
        for (PropertySource<?> source : sources) {
            collectEncryptedProperties(source, overrides);
        }

        doDecrypt(overrides);
        return overrides;
    }

    private Map<String, Object> decrypt(PropertySource<?> source) {
        Map<String, Object> overrides = new LinkedHashMap<>();
        collectEncryptedProperties(source, overrides);
        doDecrypt(overrides);
        return overrides;
    }

    private static final Pattern COLLECTION_PROPERTY = Pattern
            .compile("(\\S+)?\\[(\\d+)\\](\\.\\S+)?");

    private void collectEncryptedProperties(PropertySource<?> source,
            Map<String, Object> overrides) {

        if (source instanceof CompositePropertySource) {

            for (PropertySource<?> nested : ((CompositePropertySource) source)
                    .getPropertySources()) {
                collectEncryptedProperties(nested, overrides);
            }

        } else if (source instanceof EnumerablePropertySource) {
            Map<String, Object> otherCollectionProperties = new LinkedHashMap<>();
            boolean sourceHasDecryptedCollection = false;

            EnumerablePropertySource<?> enumerable = (EnumerablePropertySource<?>) source;
            for (String key : enumerable.getPropertyNames()) {
                Object property = source.getProperty(key);
                if (property != null) {
                    String value = property.toString();
                    if (value.startsWith("{cipher}")) {
                        overrides.put(key, value);
                        if (COLLECTION_PROPERTY.matcher(key).matches()) {
                            sourceHasDecryptedCollection = true;
                        }
                    } else if (COLLECTION_PROPERTY.matcher(key).matches()) {
                        // put non-encrypted properties so merging of index properties
                        // happens correctly
                        otherCollectionProperties.put(key, value);
                    } else {
                        overrides.remove(key);
                    }
                }
            }
            // copy all indexed properties even if not encrypted
            if (sourceHasDecryptedCollection && !otherCollectionProperties.isEmpty()) {
                overrides.putAll(otherCollectionProperties);
            }

        }
    }

    private void doDecrypt(Map<String, Object> overrides) {
        for (String key : overrides.keySet()) {
            String value = overrides.get(key).toString();
            if (value.startsWith("{cipher}")) {
                value = value.substring("{cipher}".length());
                try {
                    value = this.encryptor.decrypt(value);
                    if (logger.isDebugEnabled()) {
                        logger.debug("Decrypted: key=" + key);
                    }
                } catch (Exception e) {
                    String message = "Cannot decrypt: key=" + key;
                    if (this.failOnError) {
                        throw new IllegalStateException(message, e);
                    }
                    if (logger.isDebugEnabled()) {
                        logger.warn(message, e);
                    } else {
                        logger.warn(message);
                    }
                    // Set value to empty to avoid making a password out of the
                    // cipher text
                    value = "";
                }
                overrides.put(key, value);
            }
        }
    }
}
